{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "EheA5_j_cEwc"
      },
      "source": [
        "##### Copyright 2019 Google LLC.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "YCriMWd-pRTP"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\"); { display-mode: \"form\" }\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "OvRwFTkqcp1e"
      },
      "source": [
        "# Root Search\n",
        "\n",
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/google/tf-quant-finance/blob/master/tf_quant_finance/examples/jupyter_notebooks/Root_Search.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/google/tf-quant-finance/blob/master/tf_quant_finance/examples/jupyter_notebooks/Root_Search.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "colab_type": "code",
        "id": "uG8UAZXjeUhf",
        "colab": {}
      },
      "source": [
        "#@title Upgrade to TensorFlow 2.1+\n",
        "!pip install --upgrade tensorflow"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "uG8UAZXjeUhf"
      },
      "outputs": [],
      "source": [
        "#@title Install TF Quant Finance\n",
        "!pip install tf-quant-finance"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ku46LPe6Jswv"
      },
      "source": [
        "### This notebook demonstrates the use of low level Tensorflow Quant Finance tools for root finding using Brent's method with emphasis on the following aspects:\n",
        "\n",
        "  * **Write Once**: Tensorflow supports GPUs without any significant code changes. The same model can be run on CPU and GPU\n",
        "  * **XLA Acceleration**: The **XLA compiler** can reduce overhead associated with the Tensorflow graph by fusing operators.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "2nA2FSdTgcEM"
      },
      "outputs": [],
      "source": [
        "#@title Imports { display-mode: \"form\" }\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        " # tff for Tensorflow Finance\n",
        "import tf_quant_finance as tff \n",
        "\n",
        "root_search = tff.math.root_search\n",
        "\n",
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\",\n",
        "                        category=FutureWarning)  # suppress printing warnings"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "kK5_Zksk_gSX"
      },
      "source": [
        "### Brent's Method\n",
        "Find the risk free rate for an interest rate swap:\n",
        "$$f(x) = \\log(\\sum_{i=0}^{N} e^{-r_i t_i}) - V_{swap}$$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "13_gS0ZIBw2x",
        "outputId": "01de598e-8f48-4606-bc04-34736ed362b9"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "------------------------\n",
            "Tensorflow CPU (with auto-threading)\n",
            "('Converged:', True)\n",
            "('Estimated root:', 0.47407696)\n",
            "('Objective at root:', 1.1920929e-07)\n",
            "('Number of search steps:', 5)\n",
            "Timing:\n",
            "1000 loops, best of 3: 985 µs per loop\n",
            "------------------------\n",
            "------------------------\n",
            "Tensorflow GPU\n",
            "('Converged:', True)\n",
            "('Estimated root:', 0.47407705)\n",
            "('Objective at root:', -1.7881393e-07)\n",
            "('Number of search steps:', 5)\n",
            "Timing:\n",
            "100 loops, best of 3: 6.82 ms per loop\n",
            "------------------------\n"
          ]
        }
      ],
      "source": [
        "#@title Search Range\n",
        "\n",
        "number_of_tenors = 100  #@param\n",
        "swap_value = 0.5  #@param\n",
        "\n",
        "dtype = tf.float32\n",
        "\n",
        "tenors = range(1, number_of_tenors + 1)\n",
        "\n",
        "tf.compat.v1.reset_default_graph()\n",
        "\n",
        "t = tf.constant(tenors, dtype=dtype)\n",
        "v = tf.constant(swap_value, dtype=dtype)\n",
        "\n",
        "def valuation_fn(x):\n",
        "  return tf.reduce_logsumexp(-x * t) - v\n",
        "\n",
        "## TFF on CPU\n",
        "with tf.device(\"/cpu:0\"):\n",
        "  brent_result = root_search.brentq(valuation_fn, tf.constant(0, dtype=dtype), tf.constant(1, dtype=dtype))\n",
        "\n",
        "session = tf.Session()\n",
        "\n",
        "estimated_root, objective_at_root, num_iterations, converged = session.run(brent_result)\n",
        "\n",
        "print(\"------------------------\")\n",
        "print(\"Tensorflow CPU (with auto-threading)\")\n",
        "print(\"Converged:\", converged)\n",
        "print(\"Estimated root:\", estimated_root)\n",
        "print(\"Objective at root:\", objective_at_root)\n",
        "print(\"Number of search steps:\", num_iterations)\n",
        "print(\"Timing:\")\n",
        "%timeit session.run(brent_result)\n",
        "print(\"------------------------\")\n",
        "\n",
        "\n",
        "tf.compat.v1.reset_default_graph()\n",
        "\n",
        "t = tf.constant(tenors, dtype=dtype)\n",
        "v = tf.constant(swap_value, dtype=dtype)\n",
        "\n",
        "def valuation_fn(x):\n",
        "  return tf.reduce_logsumexp(-x * t) - v\n",
        "\n",
        "## TFF on GPU\n",
        "with tf.device(\"/gpu:0\"):\n",
        "  brent_result = root_search.brentq(valuation_fn, tf.constant(0, dtype=dtype), tf.constant(1, dtype=dtype))\n",
        "\n",
        "session = tf.Session()\n",
        "\n",
        "estimated_root, objective_at_root, num_iterations, converged = session.run(brent_result)\n",
        "\n",
        "print(\"------------------------\")\n",
        "print(\"Tensorflow GPU\")\n",
        "print(\"Converged:\", converged)\n",
        "print(\"Estimated root:\", estimated_root)\n",
        "print(\"Objective at root:\", objective_at_root)\n",
        "print(\"Number of search steps:\", num_iterations)\n",
        "print(\"Timing:\")\n",
        "%timeit session.run(brent_result)\n",
        "print(\"------------------------\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "B4Y5lQV3S9l8"
      },
      "source": [
        "### Speedup from XLA"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {
          "height": 170
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 25647,
          "status": "ok",
          "timestamp": 1563890991743,
          "user": {
            "displayName": "Cyril Chimisov",
            "photoUrl": "",
            "userId": "02803093032097482871"
          },
          "user_tz": -60
        },
        "id": "YCWMBrv-S87k",
        "outputId": "cff858d3-3287-4511-a82a-120fa54a58a2"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "------------------------\n",
            "Tensorflow CPU (compiled with XLA)\n",
            "('Converged:', True)\n",
            "('Estimated root:', 0.47407696)\n",
            "('Objective at root:', 1.1920929e-07)\n",
            "('Number of search steps:', 5)\n",
            "Timing:\n",
            "1000 loops, best of 3: 264 µs per loop\n",
            "------------------------\n",
            "------------------------\n",
            "Tensorflow GPU (compiled with XLA)\n",
            "('Converged:', True)\n",
            "('Estimated root:', 0.47407696)\n",
            "('Objective at root:', 1.1920929e-07)\n",
            "('Number of search steps:', 5)\n",
            "Timing:\n",
            "1000 loops, best of 3: 1.73 ms per loop\n",
            "------------------------\n"
          ]
        }
      ],
      "source": [
        "#@title Search Range\n",
        "\n",
        "number_of_tenors = 100  #@param\n",
        "swap_value = 0.5  #@param\n",
        "\n",
        "dtype = np.float32\n",
        "\n",
        "tenors = range(1, number_of_tenors + 1)\n",
        "\n",
        "\n",
        "tf.compat.v1.reset_default_graph()\n",
        "\n",
        "t = tf.constant(tenors, dtype=dtype)\n",
        "v = tf.constant(swap_value, dtype=dtype)\n",
        "\n",
        "def valuation_fn(x):\n",
        "  return tf.reduce_logsumexp(-x * t) - v\n",
        "\n",
        "def brent_search(objective):\n",
        "  def fn(lower_bound, upper_bound):\n",
        "    return root_search.brentq(objective, lower_bound, upper_bound)\n",
        "  return fn\n",
        "\n",
        "search_fn = brent_search(valuation_fn)\n",
        "\n",
        "## TFF on CPU compiled with XLA\n",
        "with tf.device(\"/cpu:0\"):\n",
        "  brent_result = tf.xla.experimental.compile(\n",
        "    search_fn,\n",
        "    [\n",
        "      tf.constant(0, dtype=dtype),\n",
        "      tf.constant(1, dtype=dtype),\n",
        "    ])\n",
        "\n",
        "session = tf.Session()\n",
        "\n",
        "estimated_root, objective_at_root, num_iterations, converged = session.run(brent_result)\n",
        "\n",
        "print(\"------------------------\")\n",
        "print(\"Tensorflow CPU (compiled with XLA)\")\n",
        "print(\"Converged:\", converged)\n",
        "print(\"Estimated root:\", estimated_root)\n",
        "print(\"Objective at root:\", objective_at_root)\n",
        "print(\"Number of search steps:\", num_iterations)\n",
        "print(\"Timing:\")\n",
        "%timeit session.run(brent_result)\n",
        "print(\"------------------------\")\n",
        "\n",
        "\n",
        "tf.compat.v1.reset_default_graph()\n",
        "\n",
        "t = tf.constant(tenors, dtype=dtype)\n",
        "v = tf.constant(swap_value, dtype=dtype)\n",
        "\n",
        "def valuation_fn(x):\n",
        "  return tf.reduce_logsumexp(-x * t) - v\n",
        "\n",
        "def brent_search(objective):\n",
        "  def fn(lower_bound, upper_bound):\n",
        "    return root_search.brentq(objective, lower_bound, upper_bound)\n",
        "  return fn\n",
        "\n",
        "search_fn = brent_search(valuation_fn)\n",
        "\n",
        "## TFF on GPU compiled with XLA\n",
        "with tf.device(\"/gpu:0\"):\n",
        "  brent_result = tf.xla.experimental.compile(\n",
        "    search_fn,\n",
        "    [\n",
        "      tf.constant(0, dtype=dtype),\n",
        "      tf.constant(1, dtype=dtype),\n",
        "    ])\n",
        "\n",
        "\n",
        "session = tf.Session()\n",
        "\n",
        "estimated_root, objective_at_root, num_iterations, converged = session.run(brent_result)\n",
        "\n",
        "print(\"------------------------\")\n",
        "print(\"Tensorflow GPU (compiled with XLA)\")\n",
        "print(\"Converged:\", converged)\n",
        "print(\"Estimated root:\", estimated_root)\n",
        "print(\"Objective at root:\", objective_at_root)\n",
        "print(\"Number of search steps:\", num_iterations)\n",
        "print(\"Timing:\")\n",
        "%timeit session.run(brent_result)\n",
        "print(\"------------------------\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "",
        "kind": "private"
      },
      "name": "Root_Search.ipynb",
      "provenance": [
        {
          "file_id": "1rT50234_4wb5VFNWyheBfyVlV-iuRA0t",
          "timestamp": 1553699488325
        },
        {
          "file_id": "1G2Wm6NzpdH3gPFsFTaLslgonTTkI0U0h",
          "timestamp": 1552493606270
        }
      ],
      "toc_visible": true,
      "version": "0.3.2"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
